# 首先安装必要的库
# pip install langchain langchain-community dashscope

import os
import re
from langchain.tools import tool
from langchain.prompts import PromptTemplate
from langchain.agents import create_react_agent, AgentExecutor
from langchain_community.chat_models import ChatOpenAI

# 设置通义千问API密钥
os.environ['DASHSCOPE_API_KEY'] = 'sk-4e88cf4db3e14894bafaff606d296610'

# 定义工具函数 - 使用字符串参数并手动解析
@tool
def check_market_price(flower_type: str) -> str:
    """
    查询市场上特定类型花的一般进货价格。

    参数:
    flower_type: 花的类型，如"玫瑰"、"百合"、"康乃馨"等

    返回:
    市场进货价格信息
    """
    market_prices = {
        "玫瑰": 2.5, "玫瑰花": 2.5,
        "百合": 3.8, "百合花": 3.8,
        "康乃馨": 1.2,
        "郁金香": 4.0,
        "向日葵": 2.0
    }
    price = market_prices.get(flower_type.strip('"\' '), 3.0)
    return f"{flower_type}的市场进货价格是：{price}元/支"

@tool
def calculate_price(input_str: str) -> str:
    """
    根据基础价格和加价百分比计算最终售价。
    输入格式应为: "基础价格,加价百分比" 或 "price=基础价格,percentage=加价百分比"

    参数:
    input_str: 包含基础价格和加价百分比的字符串

    返回:
    价格计算结果
    """
    try:
        # 尝试多种格式解析
        if '=' in input_str:
            # 处理 "price=2.5,percentage=5" 格式
            params = {}
            for part in input_str.split(','):
                if '=' in part:
                    key, value = part.split('=', 1)
                    params[key.strip()] = value.strip()
            base_price = float(params.get('price', params.get('base_price', 0)))
            markup_percentage = float(params.get('percentage', params.get('markup_percentage', 0)))
        else:
            # 处理 "2.5,5" 格式
            parts = input_str.split(',')
            base_price = float(parts[0].strip())
            markup_percentage = float(parts[1].strip())

        final_price = base_price * (1 + markup_percentage / 100)
        return f"基础价格: {base_price}元, 加价{markup_percentage}%后价格为: {round(final_price, 2)}元"

    except (ValueError, IndexError) as e:
        return f"计算错误: 无法解析输入参数 '{input_str}'。请使用格式: '基础价格,加价百分比'"

@tool
def get_flower_supplier_info(flower_type: str) -> str:
    """
    获取特定花材的供应商信息和建议。

    参数:
    flower_type: 花的类型

    返回:
    供应商信息和建议
    """
    supplier_info = {
        "玫瑰": "推荐供应商：云南花卉市场，质量稳定，价格优惠。建议批量采购以获得折扣。",
        "百合": "推荐供应商：广州花卉批发市场，新鲜度高，运输便捷。",
        "康乃馨": "本地供应商即可满足需求，价格波动较小。",
    }
    clean_type = flower_type.strip('"\' ')
    return supplier_info.get(clean_type, "请联系当地花卉批发市场获取最新供应商信息。")

# 工具列表
tools = [check_market_price, calculate_price, get_flower_supplier_info]

# 配置Qwen-Max模型
llm = ChatOpenAI(
    model="qwen-max",
    openai_api_key="sk-4e88cf4db3e14894bafaff606d296610",
    openai_api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
    temperature=0,
    max_tokens=2000
)

# 设置ReAct提示模板 - 明确指定输入格式
template = '''你是一个花卉价格专家，尽你所能用中文回答以下问题。如果能力不够你可以使用以下工具:

{tools}

请使用以下格式:

Question: 需要回答的输入问题
Thought: 你应该始终思考要做什么
Action: 要采取的行动，应该是 [{tool_names}] 中的一个
Action Input: 行动的输入参数（对于calculate_price，请使用"价格数字,百分比数字"格式，如"2.5,5"）
Observation: 行动的结果
... (这个 Thought/Action/Action Input/Observation 可以重复N次)
Thought: 我现在知道最终答案了
Final Answer: 对原始输入问题的最终答案

开始!

Question: {input}
Thought:{agent_scratchpad}'''

prompt = PromptTemplate.from_template(template)

# 创建ReAct Agent
agent = create_react_agent(llm, tools, prompt)

# 构建Agent执行器
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    handle_parsing_errors=True,
    verbose=True,
    max_iterations=5,
    return_intermediate_steps=True
)

# 测试工具函数
def test_tools():
    """测试工具函数是否正常工作"""
    print("测试工具函数:")
    print("1. check_market_price:", check_market_price.invoke("玫瑰"))
    print("2. calculate_price:", calculate_price.invoke("2.5,5"))
    print("3. calculate_price:", calculate_price.invoke("price=3.0,percentage=10"))
    print("4. get_flower_supplier_info:", get_flower_supplier_info.invoke("百合"))
    print()

# 自定义调用函数
def run_flower_price_query(question):
    """运行花卉价格查询"""
    try:
        print(f"正在处理问题: {question}")
        result = agent_executor.invoke({"input": question})

        if 'intermediate_steps' in result:
            print("\n执行过程详情:")
            for i, step in enumerate(result['intermediate_steps']):
                action, observation = step
                print(f"步骤 {i+1}:")
                print(f"  动作: {action.tool}")
                print(f"  输入: {action.tool_input}")
                print(f"  结果: {observation}")
                print()

        return result
    except Exception as e:
        print(f"执行错误: {str(e)}")
        return {"output": f"抱歉，处理问题时出现错误: {str(e)}"}

# 手动执行函数（备选方案）
def manual_execution():
    """手动执行流程"""
    print("\n手动执行流程:")
    print("1. 查询玫瑰价格:")
    rose_price = check_market_price.invoke("玫瑰")
    print(rose_price)

    # 从价格信息中提取数字
    price_match = re.search(r'(\d+\.?\d*)', rose_price)
    if price_match:
        base_price = float(price_match.group(1))
        print(f"2. 计算加价5%:")
        result = calculate_price.invoke(f"{base_price},5")
        print(result)
    else:
        print("无法提取价格数字")

# 示例使用
if __name__ == "__main__":
    print("=" * 60)
    print("花卉价格查询系统 - Qwen-Max ReAct Agent")
    print("=" * 60)

    # 先测试工具函数
    test_tools()

    # 测试简单问题
    simple_questions = [
        "玫瑰的市场进货价格是多少？",
        "给我一些关于百合采购的供应商建议",
    ]

    for i, question in enumerate(simple_questions, 1):
        print(f"\n简单测试 {i}: {question}")
        print("-" * 50)
        result = run_flower_price_query(question)
        if isinstance(result, dict) and 'output' in result:
            print(f"回答: {result['output']}")
        print("-" * 50)

    # 测试复杂问题
    print("\n" + "=" * 60)
    print("复杂问题测试")
    print("=" * 60)

    complex_questions = [
        "目前市场上玫瑰花的一般进货价格是多少？如果我在此基础上加价5%，应该如何定价？",
        "先查百合的价格，然后计算加价10%后的价格",
    ]

    for i, question in enumerate(complex_questions, 1):
        print(f"\n复杂测试 {i}: {question}")
        print("-" * 50)
        result = run_flower_price_query(question)
        if isinstance(result, dict) and 'output' in result:
            print(f"回答: {result['output']}")
        print("-" * 50)

    # 手动执行作为备选方案
    manual_execution()